import os
import uuid

import requests
import json
from PIL import Image
from io import BytesIO
from time import sleep
from tqdm import tqdm

# === 配置 ===
PEXELS_API_KEY = "GKHFnnRSHXC7jFEv2gjUvRV2jMFMP929T1lwHNEJ6dzfCxzl2Z9KCA5s"  # <<< 填自己的API Key
BASE_URL = "https://api.pexels.com/v1/search"
OUTPUT_DIR = "pexels_avatars"
USER_JSON = "pexels_users.json"
IMG_SIZE = (640, 640)

# 目标数量
TOTAL_IMAGES = 10
CATEGORIES = {
    "young": {"keywords": ["Japan beautiful woman"], "count": int(TOTAL_IMAGES * 1)},
   # "cartoon": {"keywords": ["cartoon"], "count": int(TOTAL_IMAGES * 0.4)},

    #"abstract": {"keywords": ["abstract"], "count": TOTAL_IMAGES - int(TOTAL_IMAGES * 0.2) - int(TOTAL_IMAGES * 0.4)},
}

HEADERS = {
    "Authorization": PEXELS_API_KEY
}

def fetch_images(keyword, per_page=10, page=10):
    params = {
        "query": keyword,
        "per_page": per_page,
        "page": page,
        "orientation": "square"
    }
    response = requests.get(BASE_URL, headers=HEADERS, params=params, timeout=10)
    if response.status_code == 200:
        return response.json().get("photos", [])
    else:
        print(f"[!] Fetch failed: {response.status_code}")
        return []


def download_and_resize(url, save_path, max_size=640):
    try:
        response = requests.get(url, timeout=10)
        if response.status_code == 200:
            img = Image.open(BytesIO(response.content)).convert("RGB")

            # 只按比例缩放，最长边等于max_size
            w, h = img.size
            if w > h:
                new_w = max_size
                new_h = int(h * (max_size / w))
            else:
                new_h = max_size
                new_w = int(w * (max_size / h))

            img = img.resize((new_w, new_h), Image.LANCZOS)

            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            img.save(save_path, quality=95)
            return True
    except Exception as e:
        print(f"[!] Download error: {e}")
    return False


def determine_gender(description):
    """简单根据描述猜性别，不精准但够用"""
    if description:
        desc = description.lower()
        if any(word in desc for word in ["woman", "female", "girl", "lady", "women"]):
            return "female"
        if any(word in desc for word in ["man", "male", "boy", "gentleman", "men"]):
            return "male"
    return "unknown"

def main():
    all_users = []
    print("[*] 开始采集图片...")

    gender_counters = {"male": 0, "female": 0, "abstract": 0}

    for category, info in CATEGORIES.items():
        count_needed = info["count"]
        keywords = info["keywords"]
        saved = 0
        page = 11

        print(f"[*] 正在下载分类 [{category}]，目标 {count_needed} 张")

        pbar = tqdm(total=count_needed)

        while saved < count_needed:
            for keyword in keywords:
                photos = fetch_images(keyword, page=page)

                if not photos:
                    print(f"[!] No more photos found for {keyword}.")
                    break

                for photo in photos:
                    if saved >= count_needed:
                        break

                    img_url = photo["src"]["large"]
                    photographer = photo.get("photographer", "Unknown")
                    description = photo.get("alt", "")

                    if category == "abstract":
                        folder = "abstract"
                    else:
                        gender = determine_gender(description)
                        if gender == "male":
                            folder = "male"
                        elif gender == "female":
                            folder = "female"
                        else:
                            folder = "male" if gender_counters["male"] <= gender_counters["female"] else "female"  # 平衡分配

                    gender_counters[folder] += 1

                    save_folder = os.path.join(OUTPUT_DIR, folder)
                    save_name = f"{uuid.uuid4().hex}.jpg"
                    save_path = os.path.join(save_folder, save_name)

                    success = download_and_resize(img_url, save_path)
                    if success:
                        all_users.append({
                            "category": category,
                            "gender_folder": folder,
                            "file_path": save_path.replace("\\", "/"),
                            "photographer": photographer,
                            "description": description
                        })
                        saved += 1
                        pbar.update(1)

                page += 1
                sleep(0.5)

        pbar.close()

    # 保存JSON
    with open(USER_JSON, "w", encoding="utf-8") as f:
        json.dump(all_users, f, indent=2, ensure_ascii=False)

    print(f"[✔] 全部完成！总共下载 {len(all_users)} 张图片，数据保存在 {USER_JSON}")

if __name__ == "__main__":
    main()
